import os
from tqdm import tqdm
import xml.etree.ElementTree as ET
import random

import config.Yolov4 as cfg

'''
作用：
创建自己的数据集
'''

def parse_voc_annotation(data_path, file_type, anno_path, use_difficult_bbox=False):
    classes = cfg.class_name
    img_inds_file = os.path.join(data_path, 'ImageSets', 'Main', file_type + '.txt')
    with open(img_inds_file, 'r') as f:
        lines = f.readlines()
        image_ids = [line.strip() for line in lines]
    with open(anno_path, 'a') as f:
        for image_id in tqdm(image_ids):
            image_path = os.path.join(data_path, 'JPEGImages', image_id + '.jpg')
            annotation = image_path
            label_path = os.path.join(data_path, 'Annotations', image_id + '.xml')
            root = ET.parse(label_path).getroot()
            objects = root.findall('object')
            for obj in objects:
                difficult = obj.find("difficult").text.strip()
                if (not use_difficult_bbox) and (int(difficult) == 1):  # difficult 表示是否容易识别，0表示容易，1表示困难
                    continue
                bbox = obj.find('bndbox')
                temp = obj.find("name").text.lower().strip()
                if temp in classes:
                    class_id = classes.index(temp)
                    xmin = bbox.find('xmin').text.strip()
                    ymin = bbox.find('ymin').text.strip()
                    xmax = bbox.find('xmax').text.strip()
                    ymax = bbox.find('ymax').text.strip()
                    annotation += ' ' + ','.join([xmin, ymin, xmax, ymax, str(class_id)])
                else:
                    continue
            annotation += '\n'

            # 为了去掉那些没有bboxes的选项
            if len(annotation.split(' '))==1:
                annotation = ''

            f.write(annotation)
    return len(image_ids)

def create_data_dir():
    dirs=['dataset','Annotations','ImageSets','JPEGImages']
    for i in dirs:
        if i==dirs[0]:
            if not os.path.exists(f'./{dirs[0]}'):
                os.mkdir(f'./{dirs[0]}')
        else:
            if not os.path.exists(f'./{dirs[0]}/{i}'):
                os.mkdir(f'./{dirs[0]}/{i}')
    if not os.path.exists(f'./{dirs[0]}/{dirs[2]}/Main'):
        os.mkdir(f'./{dirs[0]}/{dirs[2]}/Main')
    print('请在Annotations放入xml文件，在JPEGImages放入图片文件')

def voc2yolo4(xmlfilepath,saveBasePath,trainval_percent,train_percent):
    temp_xml = os.listdir(xmlfilepath)
    total_xml = []
    for xml in temp_xml:
        if xml.endswith(".xml"):
            total_xml.append(xml)
    num = len(total_xml)
    list = range(num)
    trv = int(num*trainval_percent)
    tr = int(trv*train_percent)
    tt = num - trv
    trainval = random.sample(list,trv)
    train = random.sample(trainval,tr)
    print("total",trv+tt)
    print("train size", tr)
    print("val size", trv-tr)
    print("test_size",tt)
    ftrainval = open(os.path.join(saveBasePath, 'trainval.txt'), 'w')
    ftest = open(os.path.join(saveBasePath, 'test.txt'), 'w')
    ftrain = open(os.path.join(saveBasePath, 'train.txt'), 'w')
    fval = open(os.path.join(saveBasePath, 'val.txt'), 'w')
    for i in list:
        name = total_xml[i][:-4] + '\n'
        if i in trainval:
            ftrainval.write(name)
            if i in train:
                ftrain.write(name)
            else:
                fval.write(name)
        else:
            ftest.write(name)
    ftrainval.close()
    ftrain.close()
    fval.close()
    ftest.close()

if __name__ == '__main__':
    # create_data_dir()
    trainval_percent = 1
    train_percent = 0.8
    xmlfilepath = '../data/Annotations'
    saveBasePath = r"../data/ImageSets/Main/"
    voc2yolo4(xmlfilepath,saveBasePath,trainval_percent,train_percent)
    VOC_path = os.path.abspath('../data')
    sets = ['train', 'val', 'test']
    for i in sets:
        temp_path = os.path.join('../data', f'{i}' + '_annotation.txt')
        if os.path.exists(temp_path):
            os.remove(temp_path)
        lens = parse_voc_annotation(VOC_path, i, temp_path, use_difficult_bbox=False)
        print("The number of images for train and test are :{} : {}".format(i, lens))